import pandas as pd
import numpy as np
import yfinance as yf
import matplotlib.pyplot as plt
import requests
import json
import time
from datetime import datetime, timedelta
import warnings
warnings.filterwarnings('ignore')

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号

class DingTalkRobot:
    """钉钉群机器人消息推送类"""
    
    def __init__(self, webhook_url):
        """
        初始化钉钉机器人
        
        Parameters:
        webhook_url (str): 钉钉机器人的Webhook地址
        """
        self.webhook_url = webhook_url
    
    def send_text_message(self, message, at_mobiles=None, is_at_all=False):
        """
        发送文本消息
        
        Parameters:
        message (str): 要发送的文本消息
        at_mobiles (list): 被@人的手机号列表
        is_at_all (bool): 是否@所有人
        """
        headers = {'Content-Type': 'application/json'}
        data = {
            "msgtype": "text",
            "text": {
                "content": message
            },
            "at": {
                "atMobiles": at_mobiles if at_mobiles else [],
                "isAtAll": is_at_all
            }
        }
        
        try:
            response = requests.post(
                self.webhook_url, 
                data=json.dumps(data), 
                headers=headers
            )
            if response.status_code == 200:
                print(f"钉钉消息发送成功: {message}")
            else:
                print(f"钉钉消息发送失败: {response.text}")
        except Exception as e:
            print(f"钉钉消息发送异常: {str(e)}")
    
    def send_markdown_message(self, title, text, at_mobiles=None, is_at_all=False):
        """
        发送Markdown格式消息
        
        Parameters:
        title (str): 消息标题
        text (str): Markdown格式的文本内容
        at_mobiles (list): 被@人的手机号列表
        is_at_all (bool): 是否@所有人
        """
        headers = {'Content-Type': 'application/json'}
        data = {
            "msgtype": "markdown",
            "markdown": {
                "title": title,
                "text": text
            },
            "at": {
                "atMobiles": at_mobiles if at_mobiles else [],
                "isAtAll": is_at_all
            }
        }
        
        try:
            response = requests.post(
                self.webhook_url, 
                data=json.dumps(data), 
                headers=headers
            )
            if response.status_code == 200:
                print(f"钉钉Markdown消息发送成功: {title}")
            else:
                print(f"钉钉Markdown消息发送失败: {response.text}")
        except Exception as e:
            print(f"钉钉Markdown消息发送异常: {str(e)}")

class TripleScreenTradingSystem:
    def __init__(self, symbol, ding_robot=None, 
                 long_term_period=50, medium_term_period=20, short_term_period=5):
        self.symbol = symbol
        self.long_term_period = long_term_period
        self.medium_term_period = medium_term_period
        self.short_term_period = short_term_period
        self.data = None
        self.ding_robot = ding_robot  # 钉钉机器人实例
        
    def fetch_data(self, period="6mo"):
        """获取股票数据"""
        try:
            ticker = yf.Ticker(self.symbol)
            self.data = ticker.history(period=period)
            if self.data.empty:
                raise ValueError("获取的数据为空")
            print(f"成功获取 {self.symbol} 的{len(self.data)}条数据")
            return self.data
        except Exception as e:
            error_msg = f"获取 {self.symbol} 数据失败: {str(e)}"
            print(error_msg)
            if self.ding_robot:
                self.ding_robot.send_text_message(error_msg)
            return None
        
    def calculate_macd(self, data, fast=12, slow=26, signal=9):
        """计算MACD指标"""
        exp1 = data['Close'].ewm(span=fast).mean()
        exp2 = data['Close'].ewm(span=slow).mean()
        macd = exp1 - exp2
        signal_line = macd.ewm(span=signal).mean()
        histogram = macd - signal_line
        return macd, signal_line, histogram
    
    def calculate_stochastic(self, data, period=14, k_period=3, d_period=3):
        """计算随机指标"""
        low_min = data['Low'].rolling(window=period).min()
        high_max = data['High'].rolling(window=period).max()
        
        k = 100 * (data['Close'] - low_min) / (high_max - low_min)
        k = k.rolling(window=k_period).mean()  # %K线
        d = k.rolling(window=d_period).mean()  # %D线
        return k, d
    
    def first_screen(self, data):
        """第一重滤网：长周期趋势分析"""
        # 使用MACD判断主要趋势
        macd, signal, histogram = self.calculate_macd(data)
        
        # 最近的价格行为
        recent_macd = macd.iloc[-1]
        recent_signal = signal.iloc[-1]
        recent_histogram = histogram.iloc[-1]
        
        # 判断趋势方向
        if recent_macd > recent_signal and recent_histogram > 0:
            trend = "上涨"
            direction = 1  # 看多
        elif recent_macd < recent_signal and recent_histogram < 0:
            trend = "下跌"
            direction = -1  # 看空
        else:
            trend = "震荡"
            direction = 0  # 中性
        
        return trend, direction, macd, signal, histogram
    
    def second_screen(self, data, trend_direction):
        """第二重滤网：中周期入场信号"""
        # 使用随机指标寻找入场点
        k, d = self.calculate_stochastic(data)
        
        recent_k = k.iloc[-1]
        recent_d = d.iloc[-1]
        prev_k = k.iloc[-2]
        prev_d = d.iloc[-2]
        
        signal_strength = 0
        signal_type = "无信号"
        
        # 根据主趋势方向寻找一致的信号
        if trend_direction == 1:  # 主趋势上涨，寻找买入机会
            if recent_k < 30 and recent_d < 30:  # 超卖区域
                if prev_k < prev_d and recent_k > recent_d:  # K线上穿D线
                    signal_strength = 1
                    signal_type = "买入"
                elif recent_k > recent_d and recent_k > 20:  # 保持在D线上方且脱离极端超卖
                    signal_strength = 0.5
                    signal_type = "潜在买入"
        
        elif trend_direction == -1:  # 主趋势下跌，寻找卖出机会
            if recent_k > 70 and recent_d > 70:  # 超买区域
                if prev_k > prev_d and recent_k < recent_d:  # K线下穿D线
                    signal_strength = -1
                    signal_type = "卖出"
                elif recent_k < recent_d and recent_k < 80:  # 保持在D线下方且脱离极端超买
                    signal_strength = -0.5
                    signal_type = "潜在卖出"
        
        return signal_strength, signal_type, k, d
    
    def third_screen(self, data, trend_direction, second_signal):
        """第三重滤网：短周期精确入场"""
        if second_signal == 0:  # 第二滤网无信号，直接返回
            return "无信号", 0, None, None
        
        # 使用更短周期的价格行为确认入场
        recent_price = data['Close'].iloc[-1]
        prev_price = data['Close'].iloc[-2]
        
        # 计算短期动量
        price_change = recent_price - prev_price
        percent_change = (price_change / prev_price) * 100
        
        # 根据趋势方向确认信号
        if trend_direction == 1 and second_signal > 0:  # 多头信号确认
            if price_change > 0:  # 价格上涨，确认买入
                entry_signal = "买入确认"
                strength = second_signal + 0.2  # 增强信号强度
            else:
                entry_signal = "等待更好时机"
                strength = max(0, second_signal - 0.2)  # 减弱信号强度
        
        elif trend_direction == -1 and second_signal < 0:  # 空头信号确认
            if price_change < 0:  # 价格下跌，确认卖出
                entry_signal = "卖出确认"
                strength = second_signal - 0.2  # 增强信号强度(负值)
            else:
                entry_signal = "等待更好时机"
                strength = min(0, second_signal + 0.2)  # 减弱信号强度
        
        else:
            entry_signal = "信号不一致"
            strength = 0
        
        return entry_signal, strength, price_change, percent_change
    
    def generate_signals(self):
        """生成交易信号"""
        if self.data is None:
            print("请先获取数据")
            return None
        
        # 第一重滤网：长周期趋势分析
        long_term_data = self.data
        trend, trend_direction, macd, signal, histogram = self.first_screen(long_term_data)
        
        # 第二重滤网：中周期入场信号
        medium_term_data = self.data.tail(self.medium_term_period * 3)  # 使用更多数据确保指标计算准确
        second_signal, signal_type, k, d = self.second_screen(medium_term_data, trend_direction)
        
        # 第三重滤网：短周期精确入场
        short_term_data = self.data.tail(self.short_term_period * 2)
        entry_signal, signal_strength, price_change, percent_change = self.third_screen(
            short_term_data, trend_direction, second_signal)
        
        # 准备信号详情
        signal_details = {
            "symbol": self.symbol,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "trend": trend,
            "trend_direction": trend_direction,
            "second_signal": signal_type,
            "entry_signal": entry_signal,
            "signal_strength": signal_strength,
            "current_price": self.data['Close'].iloc[-1],
            "price_change": price_change if price_change is not None else 0,
            "percent_change": percent_change if percent_change is not None else 0,
            "macd": macd.iloc[-1],
            "signal_line": signal.iloc[-1],
            "stochastic_k": k.iloc[-1] if hasattr(k, 'iloc') else 0,
            "stochastic_d": d.iloc[-1] if hasattr(d, 'iloc') else 0
        }
        
        # 发送信号到钉钉
        self.send_signal_to_dingtalk(signal_details)
        
        return signal_details
    
    def send_signal_to_dingtalk(self, signal_details):
        """发送交易信号到钉钉"""
        if self.ding_robot is None:
            print("未配置钉钉机器人，跳过消息发送")
            return
        
        symbol = signal_details['symbol']
        trend = signal_details['trend']
        signal_type = signal_details['entry_signal']
        strength = signal_details['signal_strength']
        price = round(signal_details['current_price'], 2)
        change = round(signal_details['percent_change'], 2)
        
        # 根据信号强度决定消息格式
        if abs(strength) >= 0.8:
            title = f"🚨 强烈{signal_type}信号 - {symbol}"
        elif abs(strength) >= 0.5:
            title = f"⚠️ {signal_type}信号 - {symbol}"
        else:
            title = f"ℹ️ 交易信号提醒 - {symbol}"
        
        # 创建Markdown消息内容
        markdown_text = f"""
## {title}

**标的**: {symbol}  
**当前价格**: ${price} ({change}%)  
**主要趋势**: {trend}  
**交易信号**: {signal_type}  
**信号强度**: {strength:.2f}  

### 技术指标详情
- MACD: {round(signal_details['macd'], 4)}
- Signal: {round(signal_details['signal_line'], 4)}
- Stochastic %K: {round(signal_details['stochastic_k'], 2)}
- Stochastic %D: {round(signal_details['stochastic_d'], 2)}

**时间**: {signal_details['timestamp']}
"""
        
        # 发送Markdown消息
        self.ding_robot.send_markdown_message(title, markdown_text)
        
        # 对于强烈信号，额外发送文本消息提醒
        if abs(strength) >= 0.8:
            text_message = f"{symbol} {signal_type}信号！当前价格${price}，趋势{trend}，信号强度{strength:.2f}"
            self.ding_robot.send_text_message(text_message)
    
    def plot_analysis(self, signal_details):
        """绘制分析图表"""
        if self.data is None:
            print("没有数据可绘制")
            return
        
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10))
        fig.suptitle(f'{self.symbol} 三重滤网交易系统分析', fontsize=16)
        
        # 价格和MACD图表
        ax1.plot(self.data.index, self.data['Close'], label='收盘价', linewidth=1)
        ax1.set_ylabel('价格')
        ax1.legend(loc='upper left')
        ax1.grid(True)
        
        ax1_twin = ax1.twinx()
        macd, signal, histogram = self.calculate_macd(self.data)
        ax1_twin.plot(self.data.index, macd, label='MACD', color='orange', linewidth=1)
        ax1_twin.plot(self.data.index, signal, label='Signal', color='red', linewidth=1)
        ax1_twin.bar(self.data.index, histogram, label='Histogram', color='gray', alpha=0.3)
        ax1_twin.set_ylabel('MACD')
        ax1_twin.legend(loc='upper right')
        
        # 随机指标图表
        k, d = self.calculate_stochastic(self.data)
        ax2.plot(self.data.index, k, label='%K', color='blue', linewidth=1)
        ax2.plot(self.data.index, d, label='%D', color='red', linewidth=1)
        ax2.axhline(80, color='gray', linestyle='--', alpha=0.7)
        ax2.axhline(20, color='gray', linestyle='--', alpha=0.7)
        ax2.set_ylabel('随机指标')
        ax2.legend()
        ax2.grid(True)
        ax2.set_ylim(0, 100)
        
        # 信号强度图表
        signals = []
        dates = []
        
        # 简化版信号计算（仅用于图表展示）
        for i in range(30, len(self.data)):
            sample_data = self.data.iloc[:i]
            trend, trend_direction, _, _, _ = self.first_screen(sample_data)
            medium_data = sample_data.tail(self.medium_term_period * 3)
            signal_strength, _, _, _ = self.second_screen(medium_data, trend_direction)
            signals.append(signal_strength)
            dates.append(self.data.index[i])
        
        ax3.plot(dates, signals, label='信号强度', color='purple', linewidth=1)
        ax3.axhline(0, color='black', linestyle='-', alpha=0.5)
        ax3.set_ylabel('信号强度')
        ax3.set_xlabel('日期')
        ax3.legend()
        ax3.grid(True)
        
        plt.tight_layout()
        plt.show()

# 使用示例
if __name__ == "__main__":
    # 替换为您的钉钉机器人Webhook地址
    DINGTALK_WEBHOOK = "https://oapi.dingtalk.com/robot/send?access_token=YOUR_ACCESS_TOKEN"
    
    # 创建钉钉机器人实例
    ding_robot = DingTalkRobot(DINGTALK_WEBHOOK)
    
    # 创建三重滤网交易系统实例
    trading_system = TripleScreenTradingSystem(
        symbol="AAPL",  # 股票代码
        ding_robot=ding_robot,
        long_term_period=50,
        medium_term_period=20,
        short_term_period=5
    )
    
    # 获取数据
    data = trading_system.fetch_data(period="6mo")
    
    if data is not None:
        # 生成交易信号（会自动发送到钉钉）
        signal = trading_system.generate_signals()
        
        # 打印信号详情
        print("\n交易信号详情:")
        for key, value in signal.items():
            print(f"{key}: {value}")
        
        # 绘制分析图表
        trading_system.plot_analysis(signal)